Conversation
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
|
🤖 I'm sorry @Perseus14, but I was unable to process your request. Please see the logs for more details. |
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
This Pull Request introduces significant performance optimizations for WAN models, specifically KV caching and hardware-aware dynamic image alignment padding for next-gen TPUs. The implementation of pre-computed RoPE and KV caches outside the denoising loop is an excellent architectural improvement.
🔍 General Feedback
- KV Cache Coverage: While the KV cache and RoPE pre-computation are implemented for all models, they are currently missing from the
scan_diffusion_looppaths in the T2V 2.1, T2V 2.2, and I2V 2.1 pipelines. This prevents users of the performance-oriented scan mode from benefiting from these optimizations. I have provided suggestions to propagate these caches into the scan bodies. - Hardware-Aware Padding: The dynamic adjustment of image alignment padding (256 vs 128) based on TPU type is correctly implemented across the attention and embedding layers.
- Code Organization: Moving the concatenation of embeds and RoPE computation outside the denoising loop significantly reduces redundant computation at every step.
There was a problem hiding this comment.
This Pull Request implements important performance optimizations for WAN models, including KV caching and hardware-aware dynamic image alignment padding for TPU v6e and v7x. The pre-computation of RoPE and KV caches is a significant improvement for inference efficiency.
🔍 General Feedback
- Missing Scan Loop Updates: While the optimizations are correctly implemented for the standard and CFG-cache denoising loops, the
scan_diffusion_looppaths in the T2V 2.1, T2V 2.2, and I2V 2.1 pipelines have not been updated to use the new pre-computed KV caches and RoPE. This means that users utilizing the scan mode will not see the expected performance gains and will still incur the cost of redundant RoPE computations at each step. - Hardware-Awareness: The dynamic padding logic (256 for Trillium/Ironwood, 128 otherwise) is well-integrated into the attention and embedding modules.
- Consistency: I2V 2.2 was correctly updated to handle these changes in its scan loop, but the other pipelines were missed. Synchronizing these would ensure a consistent performance profile across all WAN model variants.
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
There was a problem hiding this comment.
This Pull Request successfully implements KV Cache optimization for all WAN models (2.1 & 2.2, T2V & I2V), providing significant inference speedups on TPU. It also introduces dynamic image alignment padding, optimizing performance for Trillium (v6e) and Ironwood (v7x) accelerators. The overall implementation is solid and well-integrated into the existing pipeline and model structures.
🔍 General Feedback
- KV Cache Integration: The propagation of the KV cache from the pipeline through the transformer and attention layers is handled correctly, including support for complex CFG and caching strategies.
- Trillium Optimizations: Dynamically adjusting padding to 256 for newer TPU generations is a great performance win.
- Critical Path Bug: Identified a
NameErrorin theWanModel's TI2V path (per_token_t=True) that needs immediate attention. - Consistency: Minor inconsistencies in RoPE dummy shape calculations across pipelines were noted but do not affect correctness.
There was a problem hiding this comment.
This second part of the review adds specific inline comments and suggestions for the core model and transformer changes.
🔍 General Feedback
- TI2V Path Fix: Corrected a
NameErrorinWanModel.__call__whenper_token_tis used. - Attention Optimizations: Suggested improvements for TPU alignment detection and sequence length handling in
FlaxWanAttention.
There was a problem hiding this comment.
Detailed inline comments and suggestions for model improvements and bug fixes.
🔍 General Feedback
- TI2V Path Fix: Corrected a
NameErrorinWanModel.__call__whenper_token_tis used. - Attention Optimizations: Suggested improvements for TPU alignment detection and sequence length handling in
FlaxWanAttention.
Perseus14
force-pushed
the
wan_kv_cache
branch
3 times, most recently
from
ec2f6bb to
8467dbb
Compare
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
There was a problem hiding this comment.
This PR implements KV caching for WAN models, providing a significant optimization for inference performance. It also introduces dynamic, hardware-aware alignment padding to ensure optimal execution on next-generation TPU architectures. The changes are largely consistent across the main model variants and pipelines, but there are some omissions in the Animate model and inconsistencies in parameter propagation that should be addressed.
🔍 General Feedback
- Optimization: The pre-computation of RoPE and KV cache is a welcome improvement that aligns with JAX/Flax best practices for efficient inference.
- Hardware-Awareness: The use of dynamic alignment (128 vs 256) based on TPU type is a great addition for performance portability across Trillium and earlier generations.
- Omission:
WanAnimateTransformer3DModelappears to have been missed in this update, despite the PR's intent to cover all WAN models. - Consistency: Ensure all
WanTransformerBlockinstantiations (especially in non-scan paths) propagate the new parameters (use_base2_exp, etc.) to avoid behavioral discrepancies.
There was a problem hiding this comment.
This PR implements KV caching for WAN models, providing a significant optimization for inference performance. It also introduces dynamic, hardware-aware alignment padding to ensure optimal execution on next-generation TPU architectures (Trillium/Ironwood). While the core implementation is solid, there are omissions in the Animate model and inconsistencies in parameter propagation across non-scan code paths.
🔍 General Feedback
- Optimization: The pre-computation of RoPE and KV cache is a significant performance improvement for long sequence generation.
- Hardware-Awareness: The dynamic alignment logic based on TPU type is excellent for ensuring optimal MXU utilization across different hardware generations.
- Omission:
WanAnimateTransformer3DModelneeds to be updated with KV cache support to fulfill the goal of supporting "all WAN models". - Consistency: Ensure all
WanTransformerBlockinstantiations propagateuse_base2_expanduse_experimental_schedulerto maintain behavioral parity between scan and non-scan paths.
|
🤖 I'm sorry @Perseus14, but I was unable to process your request. Please see the logs for more details. |
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
|
🤖 I'm sorry @Perseus14, but I was unable to process your request. Please see the logs for more details. |
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
There was a problem hiding this comment.
This Pull Request implements KV caching for WAN models and adds dynamic image alignment padding for next-gen TPUs. The implementation is technically sound, significantly improves inference efficiency by avoiding redundant text embedding projections, and includes a comprehensive unit test for verification.
🔍 General Feedback
- KV Caching Logic: The addition of
compute_kv_cachetoWanModeland its integration into the denoising loops is well-executed. - TPU Optimization: The dynamic alignment padding (up to 256 for TPU v6e/v7) is a critical optimization for performance on new hardware.
- Logging: Recommend replacing standard
printstatements withmax_loggingto maintain consistency across the codebase. - Maintainability: The fix in
pyconfig.pycorrectly aligns with Python's expected behavior for__getattr__.
There was a problem hiding this comment.
This review includes additional feedback on logging practices and consistency across the new caching implementations.
🔍 General Feedback
- Consistency: Replaced remaining
printstatements withmax_loggingto ensure all diagnostic output follows the project's standard logging infrastructure.
|
🤖 I'm sorry @Perseus14, but I was unable to process your request. Please see the logs for more details. |
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
|
🤖 I'm sorry @Perseus14, but I was unable to process your request. Please see the logs for more details. |
Perseus14
force-pushed
the
wan_kv_cache
branch
2 times, most recently
from
960afef to
99a0eb0
Compare
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
|
🤖 I'm sorry @Perseus14, but I was unable to process your request. Please see the logs for more details. |
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
There was a problem hiding this comment.
This PR successfully implements KV Cache optimization for all WAN models and the WAN VACE model, significantly reducing redundant computations during the denoising loop. It also introduces hardware-aware Dynamic Image Alignment Padding to optimize performance on next-generation TPU architectures (v6e and v7x).
🔍 General Feedback
- Efficiency: The KV cache optimization is well-integrated across all pipelines (T2V, I2V, VACE) and models. While the latency savings are modest for shorter prompts, it's a solid architectural improvement.
- Hardware Optimization: The dynamic TPU type detection for image alignment padding is a great addition for future-proofing the codebase.
- Robustness: The fix for
prompt=Nonein pipelines and theAttributeErrorfix inpyconfigimprove the overall stability and usability of the CLI. - Parity: The new
wan_kv_cache_test.pyverifies that the optimization does not introduce regressions in output quality.
| timestep_proj, | ||
| encoder_hidden_states_out, | ||
| encoder_hidden_states_image_out, | ||
| encoder_attention_mask_out, | ||
| ) = self.condition_embedder( | ||
| timestep, | ||
| encoder_hidden_states, | ||
| encoder_hidden_states_image, | ||
| encoder_attention_mask, | ||
| ) = self.condition_embedder(timestep, encoder_hidden_states, encoder_hidden_states_image) | ||
| skip_embeddings=(kv_cache is not None), | ||
| ) | ||
| timestep_proj = timestep_proj.reshape(timestep_proj.shape[0], 6, -1) | ||
|
|
||
| if encoder_hidden_states_image is not None: | ||
| encoder_hidden_states = jnp.concatenate([encoder_hidden_states_image, encoder_hidden_states], axis=1) | ||
| if encoder_attention_mask is not None: | ||
| if encoder_attention_mask is None: | ||
| encoder_attention_mask = encoder_attention_mask_out | ||
|
|
||
| if encoder_hidden_states_image_out is not None: | ||
| # Workaround: When kv_cache is used, skip_embeddings=True causes raw image (1280) | ||
| # and text (4096) embeddings to have different feature dimensions. | ||
| # We pad them with zeros to match dimensions so they can be concatenated. | ||
| # This maintains the correct total sequence length for FlaxWanAttention's internal slicing. | ||
| # Since these values are ignored when cached_kv is present, the content doesn't matter. | ||
| if kv_cache is not None and encoder_hidden_states_image_out.shape[-1] != encoder_hidden_states_out.shape[-1]: | ||
| img_dim = encoder_hidden_states_image_out.shape[-1] | ||
| text_dim = encoder_hidden_states_out.shape[-1] | ||
| if img_dim < text_dim: | ||
| pad_shape = ( | ||
| encoder_hidden_states_image_out.shape[0], | ||
| encoder_hidden_states_image_out.shape[1], | ||
| text_dim - img_dim, |
There was a problem hiding this comment.
| timestep_proj, | |
| encoder_hidden_states_out, | |
| encoder_hidden_states_image_out, | |
| encoder_attention_mask_out, | |
| ) = self.condition_embedder( | |
| timestep, | |
| encoder_hidden_states, | |
| encoder_hidden_states_image, | |
| encoder_attention_mask, | |
| ) = self.condition_embedder(timestep, encoder_hidden_states, encoder_hidden_states_image) | |
| skip_embeddings=(kv_cache is not None), | |
| ) | |
| timestep_proj = timestep_proj.reshape(timestep_proj.shape[0], 6, -1) | |
| if encoder_hidden_states_image is not None: | |
| encoder_hidden_states = jnp.concatenate([encoder_hidden_states_image, encoder_hidden_states], axis=1) | |
| if encoder_attention_mask is not None: | |
| if encoder_attention_mask is None: | |
| encoder_attention_mask = encoder_attention_mask_out | |
| if encoder_hidden_states_image_out is not None: | |
| # Workaround: When kv_cache is used, skip_embeddings=True causes raw image (1280) | |
| # and text (4096) embeddings to have different feature dimensions. | |
| # We pad them with zeros to match dimensions so they can be concatenated. | |
| # This maintains the correct total sequence length for FlaxWanAttention's internal slicing. | |
| # Since these values are ignored when cached_kv is present, the content doesn't matter. | |
| if kv_cache is not None and encoder_hidden_states_image_out.shape[-1] != encoder_hidden_states_out.shape[-1]: | |
| img_dim = encoder_hidden_states_image_out.shape[-1] | |
| text_dim = encoder_hidden_states_out.shape[-1] | |
| if img_dim < text_dim: | |
| pad_shape = ( | |
| encoder_hidden_states_image_out.shape[0], | |
| encoder_hidden_states_image_out.shape[1], | |
| text_dim - img_dim, | |
| if kv_cache is not None and encoder_hidden_states_image_out.shape[-1] != encoder_hidden_states_out.shape[-1]: | |
| img_dim = encoder_hidden_states_image_out.shape[-1] | |
| text_dim = encoder_hidden_states_out.shape[-1] | |
| if img_dim < text_dim: | |
| pad_shape = ( | |
| encoder_hidden_states_image_out.shape[0], | |
| encoder_hidden_states_image_out.shape[1], | |
| text_dim - img_dim, | |
| ) | |
| encoder_hidden_states_image_out = jnp.concatenate( | |
| [ | |
| encoder_hidden_states_image_out, | |
| jnp.zeros(pad_shape, dtype=encoder_hidden_states_image_out.dtype), | |
| ], | |
| axis=-1, | |
| ) | |
| else: | |
| pad_shape = ( | |
| encoder_hidden_states_out.shape[0], | |
| encoder_hidden_states_out.shape[1], | |
| img_dim - text_dim, | |
| ) | |
| encoder_hidden_states_out = jnp.concatenate( | |
| [encoder_hidden_states_out, jnp.zeros(pad_shape, dtype=encoder_hidden_states_out.dtype)], axis=-1 | |
| ) |
…tattr and fix jax tree map
2 participants
This Pull Request implements the KV Cache optimization for all WAN models (WAN 2.1 & 2.2, both Text-to-Video and Image-to-Video) as well as the WAN VACE model.. This optimization pre-computes the Key and Value projections for text and image embeddings before the denoising loop (since they remain constant throughout)
Additionally, this PR introduces a hardware-aware Dynamic Image Alignment Padding optimization for next-generation TPUs as well as fix minor bugs and typos.
Key Changes
1. KV Cache Optimization (WAN & VACE)
FlaxWanAttention): Modifiedattention_flax.pyto acceptcached_kv. If present, key/value projections are bypassed. Added a robustcompute_kvmethod to pre-project text (and image) states.WanModel&WanVACEModel):transformer_wan.pyto support KV cache propagation. Addedcompute_kv_cacheto precompute block-level cached keys/values.skip_embeddingsinsideWanTimeTextImageEmbeddingto bypass redundant embedding layers when using cached states.transformer_wan_vace.pyfor both VACE blocks and main blocks.wan_pipeline.pyandwan_vace_pipeline_2_1.pyto accept and propagate KV caches.use_kv_cache=True.use_kv_cache: Falseto all default.ymlconfiguration files to ensure backward compatibility.2. Dynamic Image Alignment Padding (Trillium & Ironwood Optimization)
128. While optimal for older MXU tile sizes (TPU v4, v5p/v5e), next-generation hardware like Trillium (v6e) and Ironwood (v7x) utilize largerget_tpu_type()). Bothattention_flax.pyandembeddings_flax.py(NNXWanImageEmbedding) now dynamically adjust image alignment padding:v6e) and Ironwood (7x) to perfectly match larger hardware tiles.v5pand below).3. Critical Bug Fixes
prompt=Nonehandling in Pipeline: Updatedencode_promptand__call__inwan_pipeline.pyto inferbatch_sizefromprompt_embedswhenpromptisNone, unblocking tests that pass pre-computed embeddings without raw text prompts.getattrwith defaults inpyconfig: Changed__getattr__inHyperParametersclass inpyconfig.pyto raiseAttributeErrorinstead ofValueErrorwhen a key is not found. This allows standard Pythongetattr(config, "key", default)to correctly return the default value instead of crashing.jax.tree_maptojax.tree.mapin_cast_floating_to.Detailed File Changes
Models
attention_flax.py:get_tpu_typeandTpuTypefor dynamic hardware-aware image alignment padding.cached_kvrouting insideFlaxWanAttention.__call__.compute_kvsupport for both T2V and I2V cross-attentions.transformer_wan.py:skip_embeddingsparameter insideWanTimeTextImageEmbeddingto bypass redundant text/image projections.WanTransformerBlockandWanModelto handlecached_kv/kv_cachepassing.WanModel.compute_kv_cacheto precompute block-level cached keys/values across scan and non-scan layers.transformer_wan_vace.py:compute_kvtoWanVACETransformerBlock.compute_kv_cachetoWanVACEModelto precompute block-level cached keys/values for both VACE and main branches.__call__to accept and propagatekv_cacheto both VACE and main blocks.modeling_flax_utils.py:jax.tree_maptojax.tree.mapin_cast_floating_to.embeddings_flax.py:NNXWanImageEmbeddingto dynamically align to256forv6e/7xand128otherwise, avoiding shape mismatches during cross-attention.Pipelines
wan_pipeline.py:transformer_forward_pass,transformer_forward_pass_full_cfg, andtransformer_forward_pass_cfg_cacheto accept and passkv_cache.NoneTypeerror when calculating batch size by checking ifpromptisNone.wan_vace_pipeline_2_1.py:use_kv_cacheparameter to pipeline calls.kv_cachebefore the denoising loop and passed it totransformer_forward_passinrun_inference.wan_pipeline_2_1.py&wan_pipeline_2_2.py:use_kv_cacheparameter to pipeline calls and pre-computedkv_cacheandrotary_embprior to the denoising loop.wan_pipeline_i2v_2p1.py&wan_pipeline_i2v_2p2.py:kv_cachesupport for I2V workflows.generate_wan.py:Configs
base_wan_1_3b.yml,base_wan_14b.yml,base_wan_27b.yml,base_wan_i2v_14b.yml,base_wan_i2v_27b.yml):use_kv_cache: Falsefor all default configs.promptin I2V config filesPerformance Note
Note
Test Configuration: 720p | 81 frames | 40 steps
Hardware: TPU 7x-8
JAX Version: v0.10.0
Conclusion: No visual change from baseline across all tested variants.